#!/home/san/.conda/envs/dp/bin/python
"""
@file: datasets.py
@docs: 对dicom做进一步处理，并将其保存为卷积网络的输入样本
"""


import pandas as pd
import glob
import os
import sys
import configparser
import cv2
from dcm_utils import dicom2array, dicom_metainfo


def get_info(trainPath, jsonPath):
    """
    :param trainPath: 训练所用的dicom数据根目录
    :param jsonPath: 标注文件目录
    :return: 图片路径及与之对应的标注
    """
    annotation_info = pd.DataFrame(columns=('studyUid', 'seriesUid', 'instanceUid', 'annotation'))
    json_df = pd.read_json(jsonPath)

    for idx in json_df.index:
        studyUid = json_df.loc[idx, "studyUid"]
        seriesUid = json_df.loc[idx, "train_data"][0]['seriesUid']
        instanceUid = json_df.loc[idx, "train_data"][0]['instanceUid']
        annotation = json_df.loc[idx, "train_data"][0]['annotation']
        row = pd.Series(
            {'studyUid': studyUid, 'seriesUid': seriesUid, 'instanceUid': instanceUid, 'annotation': annotation})
        annotation_info = annotation_info.append(row, ignore_index=True)
    dcm_paths = glob.glob(os.path.join(trainPath, "**", "**.dcm"))  # 具体的图片路径
    # print(len(dcm_paths))
    # 'studyUid','seriesUid','instanceUid'
    tag_list = ['0020|000d', '0020|000e', '0008|0018']
    dcm_info = pd.DataFrame(columns=('dcmPath', 'studyUid', 'seriesUid', 'instanceUid'))
    for dcm_path in dcm_paths:
        try:
            studyUid, seriesUid, instanceUid = dicom_metainfo(dcm_path, tag_list)
            row = pd.Series(
                {'dcmPath': dcm_path, 'studyUid': studyUid, 'seriesUid': seriesUid, 'instanceUid': instanceUid})
            dcm_info = dcm_info.append(row, ignore_index=True)
        except:
            continue
    # print(dcm_info)
    # 标注文件拼接
    result = pd.merge(annotation_info, dcm_info, on=['studyUid', 'seriesUid', 'instanceUid'])
    result = result.set_index('dcmPath')['annotation']  # 然后把index设置为路径，值设置为annotation
    # print(result)
    result.to_csv("../result.csv")
    return result


def generator(samples):
    """
    根据samples生成x_train
    :return: x, _y= [], []
    """
    x, _y = [], []
    for i in range(len(samples)):
        x.append(samples[i][0])
        _y.append(samples[i][1])
    return x, _y


def save_as_jpg():
    """将dcm文件保存为jpg图片"""
    train_path = '../dataset/SparkAI/train/train_data_dcm'
    json_path = '../dataset/SparkAI/train/train_label/lumbar_train51_annotation.json'
    res = get_info(train_path, json_path)

    for i in range(len(res)):
        img_dir = res.index[i]  # 获取图片的地址

        img = dicom2array(img_dir)
        img_name = img_dir.replace('/', '_')
        cv2.imwrite('../dataset/SparkAI/train/train_data/' + img_name + '.jpg', img)



CONFIG_FILE = '../config.cfg'


def main():
    # 从配置文件中获取信息
    conf = configparser.ConfigParser()
    conf.read(CONFIG_FILE)
    train_path = conf.get("FL_CONFIG", "TRAIN_PATH")
    json_path = conf.get("FL_CONFIG", "JSON_PATH")
    x_train_save_path = conf.get("FL_CONFIG", "x_train_save_path")
    y_train_save_path = conf.get("FL_CONFIG", "y_train_save_path")

    # 获取路径和标注文件的对应关系
    res = get_info(train_path, json_path)


    # for i in range(len(res)):
    #     img_dir = res.index[i]
    #     # print(img_dir)
    #     tags = res[i]
    #     print(tags)

    # sample = get_sample(res)
    # print("sample.shape: ", len(sample))
    # print(sample)
    # x_train, y_train = generator(sample)
    # x_train_save = np.reshape(x_train, (len(x_train), -1))
    #
    # np.save(x_train_save_path, x_train_save)
    # np.save(y_train_save_path, y_train)


if __name__ == '__main__':
    save_as_jpg()